// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include "zencore.h"

#include <zenbase/refcount.h>
#include <zencore/iobuffer.h>
#include <zencore/memory.h>

#include <memory.h>

namespace zen {

class SharedBuffer;

/**
 * Reference to a memory buffer with a single owner
 *
 * Internally
 */
class UniqueBuffer
{
public:
	UniqueBuffer()				 = default;
	UniqueBuffer(UniqueBuffer&&) = default;
	UniqueBuffer& operator=(UniqueBuffer&&) = default;
	UniqueBuffer(const UniqueBuffer&)		= delete;
	UniqueBuffer& operator=(const UniqueBuffer&) = delete;

	ZENCORE_API explicit UniqueBuffer(IoBufferCore* Owner);

	[[nodiscard]] void*		  GetData() { return m_Buffer ? m_Buffer->MutableDataPointer() : nullptr; }
	[[nodiscard]] const void* GetData() const { return m_Buffer ? m_Buffer->DataPointer() : nullptr; }
	[[nodiscard]] size_t	  GetSize() const { return m_Buffer ? m_Buffer->DataBytes() : 0; }

	operator MutableMemoryView() { return GetMutableView(); }
	operator MemoryView() const { return GetView(); }

	/**
	 * Returns true if this does not point to a buffer owner.
	 *
	 * A null buffer is always owned, materialized, and empty.
	 */
	[[nodiscard]] inline bool IsNull() const { return m_Buffer.IsNull(); }

	/** Reset this to null. */
	ZENCORE_API void Reset();

	[[nodiscard]] inline MutableMemoryView GetMutableView() { return MutableMemoryView(GetData(), GetSize()); }
	[[nodiscard]] inline MemoryView		   GetView() const { return MemoryView(GetData(), GetSize()); }

	/** Make an uninitialized owned buffer of the specified size. */
	[[nodiscard]] ZENCORE_API static UniqueBuffer Alloc(uint64_t Size);

	/** Make a non-owned view of the input. */
	[[nodiscard]] ZENCORE_API static UniqueBuffer MakeMutableView(void* DataPtr, uint64_t Size);

	/**
	 * Convert this to an immutable shared buffer, leaving this null.
	 *
	 * Steals the buffer owner from the unique buffer.
	 */
	[[nodiscard]] ZENCORE_API SharedBuffer MoveToShared();

private:
	// This may be null, for a default constructed UniqueBuffer only
	RefPtr<IoBufferCore> m_Buffer;

	friend class SharedBuffer;
};

/**
 * Reference to a memory buffer with shared ownership
 */
class SharedBuffer
{
public:
	SharedBuffer() = default;
	ZENCORE_API explicit SharedBuffer(UniqueBuffer&&);
	inline explicit SharedBuffer(IoBufferCore* Owner) : m_Buffer(Owner) {}
	ZENCORE_API explicit SharedBuffer(IoBuffer&& Buffer) : m_Buffer(std::move(Buffer.m_Core)) {}
	ZENCORE_API explicit SharedBuffer(const IoBuffer& Buffer) : m_Buffer(Buffer.m_Core) {}
	ZENCORE_API explicit SharedBuffer(RefPtr<IoBufferCore>&& Owner) : m_Buffer(std::move(Owner)) {}

	[[nodiscard]] const void* GetData() const
	{
		if (m_Buffer)
		{
			return m_Buffer->DataPointer();
		}
		return nullptr;
	}

	[[nodiscard]] size_t GetSize() const
	{
		if (m_Buffer)
		{
			return m_Buffer->DataBytes();
		}
		return 0;
	}

	inline void MakeImmutable()
	{
		ZEN_ASSERT(m_Buffer);
		m_Buffer->SetIsImmutable(true);
	}

	/** Returns a buffer that is owned, by cloning if not owned. */
	[[nodiscard]] ZENCORE_API SharedBuffer MakeOwned() const&;
	[[nodiscard]] ZENCORE_API SharedBuffer MakeOwned() &&;

	[[nodiscard]] bool		  IsOwned() const { return !m_Buffer || m_Buffer->IsOwned(); }
	[[nodiscard]] inline bool IsNull() const { return !m_Buffer; }
	inline void				  Reset() { m_Buffer = nullptr; }
	inline bool				  GetFileReference(IoBufferFileReference& OutRef) const
	{
		if (const IoBufferExtendedCore* Core = m_Buffer->ExtendedCore())
		{
			return Core->GetFileReference(OutRef);
		}
		else
		{
			return false;
		}
	}

	[[nodiscard]] MemoryView GetView() const
	{
		if (m_Buffer)
		{
			return MemoryView(m_Buffer->DataPointer(), m_Buffer->DataBytes());
		}
		else
		{
			return MemoryView();
		}
	}
	operator MemoryView() const { return GetView(); }

	/** Returns true if this points to a buffer owner. */
	[[nodiscard]] inline explicit operator bool() const { return !IsNull(); }

	[[nodiscard]] inline IoBuffer AsIoBuffer() const { return IoBuffer(m_Buffer); }

	SharedBuffer& operator=(UniqueBuffer&& Rhs)
	{
		m_Buffer = std::move(Rhs.m_Buffer);
		return *this;
	}

	std::strong_ordering operator<=>(const SharedBuffer& Rhs) const = default;

	/** Make a non-owned view of the input */
	[[nodiscard]] inline static SharedBuffer MakeView(MemoryView View) { return MakeView(View.GetData(), View.GetSize()); }
	/** Make a non-owning view of the memory of the contiguous container. */
	[[nodiscard]] inline static SharedBuffer MakeView(const ContiguousRange auto& Container)
	{
		std::span Span = Container;
		return MakeView(Span.data(), Span.size() * sizeof(typename decltype(Span)::element_type));
	}
	/** Make a non-owned view of the input */
	[[nodiscard]] ZENCORE_API static SharedBuffer MakeView(const void* Data, uint64_t Size);
	/** Make a non-owned view of the input */
	[[nodiscard]] ZENCORE_API static SharedBuffer MakeView(MemoryView View, SharedBuffer OuterBuffer);
	/** Make an owned clone of the buffer */
	[[nodiscard]] ZENCORE_API SharedBuffer Clone();
	/** Make an owned clone of the memory in the input view */
	[[nodiscard]] ZENCORE_API static SharedBuffer Clone(MemoryView View);

private:
	RefPtr<IoBufferCore> m_Buffer;
};

void sharedbuffer_forcelink();

}  // namespace zen
